
// 拟合直线

#include <ros/ros.h>
#include <nav_msgs/Path.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud2_iterator.h>

#include <iostream>
#include <random>
using namespace std;

class LineTest{
    public:
        LineTest();

        // 发布点云数据
        void PublishPointCloud2(vector<geometry_msgs::Point>);
        // 发布直线数据
        void PublishLine(geometry_msgs::Point start, geometry_msgs::Point end);
        // 梯度下降法拟合直线 y = k*x + b
        void fitLineGradientDescent(vector<geometry_msgs::Point> points_in, double& k, double& b);
        // 最小二乘法拟合直线 y = k*x + b
        // https://blog.csdn.net/xinjiang666/article/details/103782544
        // https://blog.csdn.net/stf1065716904/article/details/107594710
        void fitLineLeastSquare(vector<geometry_msgs::Point> points_in, double& k, double& b);

    private:
        ros::NodeHandle nh_;
        ros::Publisher linePub_;
        ros::Publisher pointsPub_;
        vector<geometry_msgs::Point> points__;
};


LineTest::LineTest(){
    linePub_ = nh_.advertise<nav_msgs::Path>("/fitting_line", 1, true);
    pointsPub_ = nh_.advertise<sensor_msgs::PointCloud2>("/fitting_points", 1, true);

    ros::Duration(1.0).sleep();

    // 生成一组随机数
    random_device rd;
    mt19937 gen(rd());
    uniform_real_distribution<> dis_x(0.5, 1.0);
    uniform_real_distribution<> dis_y(-5.0, 5.0);
    cout << "生成随机数据" << endl;
    for (int i = 0; i < 50; ++i) {
        geometry_msgs::Point point;
        point.y = dis_y(gen);
        point.x = dis_x(gen) * point.y;
        cout << "--> [" << i << "] data: x= " << point.x << ", y=" << point.y << endl;
        points__.emplace_back(point);
    }

    // 发布点云数据
    PublishPointCloud2(points__);

    // 拟合直线
    double k = 0.0;
    double b = 0.0;
    fitLineGradientDescent(points__, k, b);   //梯度下降法拟合直线
    fitLineLeastSquare(points__, k, b);    //最小二乘法拟合直线

    // 发布拟合的直线
    std::cout << "拟合的直线: y = " << k << "x + " << b << std::endl; 
    // 计算直线方程
    geometry_msgs::Point start;
    start.x = -5.0;
    start.y = k * start.x + b;
    cout << "start: x=" << start.x << ", y=" << start.y << endl;
    geometry_msgs::Point end;
    end.x = 5.0;
    end.y = k * end.x + b;
    cout << "end: x=" << end.x << ", y=" << end.y << endl;
    // 发布拟合直线
    PublishLine(start, end);
}


/******************************************************
 * 拟合直线
 ******************************************************/
// 梯度下降法拟合直线
// https://zhuanlan.zhihu.com/p/63799123
void LineTest::fitLineGradientDescent(vector<geometry_msgs::Point> points_in, double& k, double& b){
    int n = points_in.size();
    int max_iter = 10000;           // 最大迭代次数
    double learning_rate = 0.01;    // 学习率

    for(int i=0; i<max_iter; ++i){
        double da = 0.0;
        double db = 0.0;  

        // 计算梯度
        for(int j = 0; j < n; ++j) {
            auto point = points_in[j];
            double prediction = k * point.x + b;
            da += -2 * point.x * (point.y - prediction);
            db += -2 * (point.y - prediction);
        }  

        // 更新参数
        k -= (learning_rate / n) * da;
        b -= (learning_rate / n) * db;
    }
    std::cout << "[fitLineGradientDescent] y = " << k << "x + " << b << std::endl; 
}


// 最小二乘法拟合直线
void LineTest::fitLineLeastSquare(vector<geometry_msgs::Point> points_in, double& k, double& b){
    int n = points_in.size();
	double xx_sum = 0;
	double x_sum = 0;
	double y_sum = 0;
	double xy_sum = 0;
	for (int i = 0; i < n; i++){
        auto point = points_in[i];
		x_sum += point.x; //x的累加和
		y_sum += point.y; //y的累加和
		xx_sum += point.x * point.x; //x的平方累加和
		xy_sum += point.x * point.y; //x，y的累加和
	}
	k = (n*xy_sum - x_sum * y_sum) / (n*xx_sum - x_sum * x_sum); //根据公式求解k
	b = (-x_sum * xy_sum + xx_sum*y_sum) / (n*xx_sum - x_sum * x_sum);//根据公式求解b

    std::cout << "[fitLineLeastSquare] y = " << k << "x + " << b << std::endl; 
}

/******************************************************
 * 数据发布 
 ******************************************************/

// 发布点云数据
void LineTest::PublishPointCloud2(vector<geometry_msgs::Point> points){
    sensor_msgs::PointCloud2 cloud_msg;
    cloud_msg.header.frame_id = "map";
    cloud_msg.header.stamp = ros::Time::now();
    cloud_msg.height = 1;
    cloud_msg.width = points.size();
    cloud_msg.is_dense = false;

    sensor_msgs::PointCloud2Modifier modifier(cloud_msg);
    modifier.setPointCloud2FieldsByString(1, "xyz");
    modifier.resize(points.size());

    // 使用迭代器填充点云消息的数据
    sensor_msgs::PointCloud2Iterator<float> iter_x(cloud_msg, "x");
    sensor_msgs::PointCloud2Iterator<float> iter_y(cloud_msg, "y");
    sensor_msgs::PointCloud2Iterator<float> iter_z(cloud_msg, "z");
    for(const auto& point : points) {
        *iter_x = point.x;
        *iter_y = point.y;
        *iter_z = point.z;

        ++iter_x;
        ++iter_y;
        ++iter_z;
    }
    pointsPub_.publish(cloud_msg);
}


// 发布直线数据
void LineTest::PublishLine(geometry_msgs::Point start, geometry_msgs::Point end){
    nav_msgs::Path path_msg;
    path_msg.header.frame_id = "map";
    path_msg.header.stamp = ros::Time::now();
    path_msg.poses.resize(2);
    path_msg.poses.front().pose.position = start;
    path_msg.poses.back().pose.position = end;
    linePub_.publish(path_msg);
}


/******************************************************
 * 主函数 * 
 ******************************************************/
int main(int argc, char** argv){
    ros::init(argc, argv, "test_line_gradient");
    setlocale(LC_ALL,"");   //ROS_INFO 中文显示, 需要添加头文件#include <sstream>
	
    ROS_INFO("test_line_gradient initializating......");
    ros::NodeHandle nh;
    LineTest line_fitting;
    ros::spin();
    return 0;
}